#!/usr/bin/env python3

import argparse
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from scipy.optimize import curve_fit

parser = argparse.ArgumentParser(description='Do observation error statistics.')
parser.add_argument('--synop', action='store_true')
parser.add_argument('--raob' , action='store_true')
parser.add_argument('--amdar' , action='store_true')
parser.add_argument('--profiler' , action='store_true')
parser.add_argument('--omb2', help='HDF5 file containing OMB-OMB results from previous step', required=True)
parser.add_argument('--var', help='Variable to do statistics', required=True)
parser.add_argument('--obs-p', help='Observation pressure (hPa) for RAOB, amdar', type=float)
parser.add_argument('--bin-size', required=True, type=float)
parser.add_argument('--max-dist', required=True, type=float)
parser.add_argument('--output', required=True)
args = parser.parse_args()

df = pd.read_hdf(args.omb2)

if args.raob:
	df = df.loc[df['p'] == args.obs_p]

bins = np.arange(-0.01, args.max_dist, args.bin_size)

df = df.assign(bined=lambda row: pd.cut(row.d_km, bins))

df = df.groupby(['bined']).agg(mean_omb_var=(f'omb_{args.var}_var', np.nanmean))
print(df)

bins[0] = 0

def fit_func(x, A, B):
	return A * np.exp(-(x / B)**2)

y = np.array(df.mean_omb_var)
x = np.linspace(0, len(y) - 1, len(y))
popt, pcov = curve_fit(fit_func, x[1:], y[1:])

obs_error = np.sqrt(y[0] - fit_func(0, *popt))
print(f'*** Estimated observation error is {obs_error} ***')

ax1 = plt.bar(bins[:-1], y, width=args.bin_size, label='bin')
ax2 = plt.plot(bins[:-1], fit_func(x, *popt), 'r-', label='fit')
plt.title(f'RAOB ({args.obs_p}hPa)')
plt.xlabel('Separation (KM)')
plt.ylabel(f'omb_{args.var}_var')

plt.savefig(args.output)
